#include "x64.h"

NTSTATUS
CurrentProcessVaToPa(VA va, PPA pa)
{
	CR3 Cr3 = { .Value = __readcr3() };

	PA     L1PtPa = { .Fields4KB.PPN = Cr3.Fields.PPN };
	PL1PTE L1Pt   = MmGetVirtualForPhysical(L1PtPa.AsLargeInteger);
	L1PTE  L1Pte  = L1Pt[va.Fields.VPN1];
	if (L1Pte.Fields.P == 0) {
		return STATUS_INVALID_ADDRESS;
	}

	PA     L2PtPa = { .Fields4KB.PPN = L1Pte.Fields.PPN };
	PL2PTE L2Pt   = MmGetVirtualForPhysical(L2PtPa.AsLargeInteger);
	L2PTE  L2Pte  = L2Pt[va.Fields.VPN2];
	if (L2Pte.Fields.P == 0) {
		return STATUS_INVALID_ADDRESS;
	}
	if (L2Pte.Fields.PS == 1) {
		*pa = (PA){ .Fields1GB = {
			.PPN = L2Pte.Fields1GB.PPN,
			.PPO = va.Fields1GB.VPO} };
		return STATUS_SUCCESS;
	}

	PA     L3PtPa = { .Fields4KB.PPN = L2Pte.Fields.PPN };
	PL3PTE L3Pt   = MmGetVirtualForPhysical(L3PtPa.AsLargeInteger);
	L3PTE  L3Pte  = L3Pt[va.Fields.VPN3];
	if (L3Pte.Fields.P == 0) {
		return STATUS_INVALID_ADDRESS;
	}
	if (L3Pte.Fields.PS == 1) {
		*pa = (PA){ .Fields2MB = {
			.PPN = L3Pte.Fields2MB.PPN,
			.PPO = va.Fields2MB.VPO } };
		return STATUS_SUCCESS;
	}

	PA L4PtPa = { .Fields4KB.PPN = L3Pte.Fields.PPN };
	PL4PTE L4Pt = MmGetVirtualForPhysical(L4PtPa.AsLargeInteger);
	L4PTE L4Pte = L4Pt[va.Fields.VPN4];
	if (L4Pte.Fields.P == 0) {
		return STATUS_INVALID_ADDRESS;
	}

	*pa = (PA){ .Fields4KB = {
		.PPN = L4Pte.Fields.PPN,
		.PPO = va.Fields4KB.VPO } };

	return STATUS_SUCCESS;
}
NTSTATUS
VaToPa(PEPROCESS Process, VA Va, PPA Pa)
{
	if (PsGetCurrentProcess() == Process) return CurrentProcessVaToPa(Va, Pa);

	NTSTATUS Status;
	KAPC_STATE ApcState;
	KeStackAttachProcess(Process, &ApcState);
	Status = CurrentProcessVaToPa(Va, Pa);
	KeUnstackDetachProcess(&ApcState);
	return Status;
}
